/*
 * @Author: LVGRAPE
 * @LastEditors: LVGRAPE
 */
/*
 * Copyright (c) 2006-2021, RT-Thread Development Team
 *
 * SPDX-License-Identifier: Apache-2.0
 *
 * Change Logs:
 * Date           Author       Notes
 * 2023-09-07     UAVI02       the first version
 */
#include "zino.h"
#include "imu.h"
#include "fc_sensors.h"
#include <math.h>
#include "maths.h"
#include "fc_state.h"

#define DBG_TAG "imu"
#define DBG_LVL DBG_LOG
#include <rtdbg.h>

#define USELPF
#define IMU_UPDATE_DT 0.01

ZINO_CONFIG_REGISTER(imu_offset_t, imu_offset, 1, 1);

static float invSqrt(float x)
{
    // return 1.0f / sqrtf(x);

    float halfx = 0.5f * x;
    float y = x;
    long i = *(long *)&y;             // get bits for floating value
    i = 0x5f3759df - (i >> 1);        // gives initial guess
    y = *(float *)&i;                 // convert bits back to float
    y = y * (1.5f - (halfx * y * y)); // Newton step
    return y;
}

// 计算旋转矩阵
static void imuComputeRotationMatrix(m_imu_t *imu)
{

    /**
     * @brief R = [
     *  [1-2(q2²+q3²)   2(q1q2-q0q3)   2(q1q3+q0q2) ]
     *  [2(q1q2+q0q3)   1-2(q1²+q3²)   2(q2q3-q0q1) ]
     *  [2(q1q3-q0q2)   2(q2q3+q0q1)   1-2(q1²+q2²) ]
     * ]
     *
     */
    // 计算q1的平方
    float q1q1 = imu->q1 * imu->q1;
    // 计算q2的平方
    float q2q2 = imu->q2 * imu->q2;
    // 计算q3的平方
    float q3q3 = imu->q3 * imu->q3;

    // 计算q0和q1的乘积
    float q0q1 = imu->q0 * imu->q1;
    // 计算q0和q2的乘积
    float q0q2 = imu->q0 * imu->q2;
    // 计算q0和q3的乘积
    float q0q3 = imu->q0 * imu->q3;
    // 计算q1和q2的乘积
    float q1q2 = imu->q1 * imu->q2;
    // 计算q1和q3的乘积
    float q1q3 = imu->q1 * imu->q3;
    // 计算q2和q3的乘积
    float q2q3 = imu->q2 * imu->q3;

    // 计算旋转矩阵的第一行
    imu->rMat[0][0] = 1.0f - 2.0f * q2q2 - 2.0f * q3q3;
    imu->rMat[0][1] = 2.0f * (q1q2 + -q0q3);
    imu->rMat[0][2] = 2.0f * (q1q3 - -q0q2);

    // 计算旋转矩阵的第二行
    imu->rMat[1][0] = 2.0f * (q1q2 - -q0q3);
    imu->rMat[1][1] = 1.0f - 2.0f * q1q1 - 2.0f * q3q3;
    imu->rMat[1][2] = 2.0f * (q2q3 + -q0q1);

    // 计算旋转矩阵的第三行
    imu->rMat[2][0] = 2.0f * (q1q3 + -q0q2);
    imu->rMat[2][1] = 2.0f * (q2q3 - -q0q1);
    imu->rMat[2][2] = 1.0f - 2.0f * q1q1 - 2.0f * q2q2;
}

static void imuTransformVectorBodyToEarth(m_imu_t *imu, fc_sensor_data_t *data, Axis_3_float_t *acc)
{

    const float x = imu->rMat[0][0] * data->acc.acclpf.x + imu->rMat[0][1] * data->acc.acclpf.y + imu->rMat[0][2] * data->acc.acclpf.z;
    const float y = imu->rMat[1][0] * data->acc.acclpf.x + imu->rMat[1][1] * data->acc.acclpf.y + imu->rMat[1][2] * data->acc.acclpf.z;
    const float z = imu->rMat[2][0] * data->acc.acclpf.x + imu->rMat[2][1] * data->acc.acclpf.y + imu->rMat[2][2] * data->acc.acclpf.z;

    float yawRad = atan2f(imu->rMat[1][0], imu->rMat[0][0]);
    float cosy = cosf(yawRad);
    float siny = sinf(yawRad);
    float vx = x * cosy + y * siny;
    float vy = y * cosy - x * siny;

    // Y towards left, X towards front, Z towards up
    acc->x = vx;
    acc->y = vy;
    acc->z = z - 1.0f;

    // rt_printf_float("acc->x ", acc->x, " ");
    // rt_printf_float("acc->y ", acc->y, " ");
    // rt_printf_float("acc->z ", acc->z, " \n");
}

#define SPIN_RATE_LIMIT 20

void MahonyAHRS(m_imu_t *imu, fc_sensor_data_t *data, float dt)
{
    float ex = 0, ey = 0, ez = 0;

    float ax = data->acc.accSum.x / data->acc.sumCount.x;
    float ay = data->acc.accSum.y / data->acc.sumCount.y;
    float az = data->acc.accSum.z / data->acc.sumCount.z;
    float gx = DEGREES_TO_RADIANS(data->gyro.gyroSum.x / data->gyro.sumCount.x);
    float gy = DEGREES_TO_RADIANS(data->gyro.gyroSum.y / data->gyro.sumCount.y);
    float gz = DEGREES_TO_RADIANS(data->gyro.gyroSum.z / data->gyro.sumCount.z);

    data->acc.accSum.x = 0;
    data->acc.accSum.y = 0;
    data->acc.accSum.z = 0;
    data->gyro.gyroSum.x = 0;
    data->gyro.gyroSum.y = 0;
    data->gyro.gyroSum.z = 0;

    data->acc.sumCount.x = 0;
    data->acc.sumCount.y = 0;
    data->acc.sumCount.z = 0;
    data->gyro.sumCount.x = 0;
    data->gyro.sumCount.y = 0;
    data->gyro.sumCount.z = 0;

    // return;

    // Calculate general spin rate (rad/s)
    const float spin_rate = sqrtf(sq(gx) + sq(gy) + sq(gz));

    float recipNorm = sq(ax) + sq(ay) + sq(az);

    /*单位化加速计测量值*/
    recipNorm = invSqrt(recipNorm);
    ax *= recipNorm;
    ay *= recipNorm;
    az *= recipNorm;
    /*加速计读取的方向与重力加速计方向的差值，用向量叉乘计算*/
    ex = (ay * imu->rMat[2][2] - az * imu->rMat[2][1]);
    ey = (az * imu->rMat[2][0] - ax * imu->rMat[2][2]);
    ez = (ax * imu->rMat[2][1] - ay * imu->rMat[2][0]);
    /*误差累计，与积分常数相乘*/
    if (imu->ki > 0.0f)
    {
        // imu->integralFBx += imu->ki * ex * dt;
        // imu->integralFBy += imu->ki * ey * dt;
        // imu->integralFBz += imu->ki * ez * dt;

        // Stop integrating if spinning beyond the certain limit
        if (spin_rate < DEGREES_TO_RADIANS(SPIN_RATE_LIMIT))
        {
            imu->integralFBx += imu->ki * ex * dt;
            imu->integralFBy += imu->ki * ey * dt;
            imu->integralFBz += imu->ki * ez * dt;
        }

        if (FC_STATE_CHECK(FC_IMU_CALIBRATING))
        {

            imu->integralFBx = 0;
            imu->integralFBy = 0;
            imu->integralFBz = 0;
        }
    }
    else
    {
        imu->integralFBx = 0.0f;
        imu->integralFBy = 0.0f;
        imu->integralFBz = 0.0f;
    }
    /*用叉积误差来做PI修正陀螺零偏，即抵消陀螺读数中的偏移量*/
    float kp = (millis() < 500) ? imu->kp * 10 : imu->kp;
    gx += kp * ex + imu->integralFBx;
    gy += kp * ey + imu->integralFBy;
    gz += kp * ez + imu->integralFBz;
    /* 一阶近似算法，四元数运动学方程的离散化形式和积分 */
    float halfT = 0.5f * dt;
    const float qa = imu->q0; // last q
    const float qb = imu->q1; // last q
    const float qc = imu->q2; // last q
    const float qd = imu->q3; // last q

    imu->q0 += (-qb * gx - qc * gy - qd * gz) * halfT;
    imu->q1 += (qa * gx + qc * gz - qd * gy) * halfT;
    imu->q2 += (qa * gy - qb * gz + qd * gx) * halfT;
    imu->q3 += (qa * gz + qb * gy - qc * gx) * halfT;
    /*单位化四元数*/
    recipNorm = invSqrt(sq(imu->q0) + sq(imu->q1) + sq(imu->q2) + sq(imu->q3));
    imu->q0 = imu->q0 * recipNorm;
    imu->q1 = imu->q1 * recipNorm;
    imu->q2 = imu->q2 * recipNorm;
    imu->q3 = imu->q3 * recipNorm;
    /*计算旋转矩阵*/
    imuComputeRotationMatrix(imu);
    /*计算roll pitch yaw 欧拉角*/
    data->attitude.roll = RADIANS_TO_DEGREES(atan2f(imu->rMat[2][1], imu->rMat[2][2]));
    data->attitude.pitch = RADIANS_TO_DEGREES(-asinf(imu->rMat[2][0]));
    data->attitude.yaw = RADIANS_TO_DEGREES(atan2f(imu->rMat[1][0], imu->rMat[0][0]));
    Get_imu_data(data, &zdrone_state);
    imuTransformVectorBodyToEarth(imu, data, &zdrone_state.acc);

    zdrone_state.attitude.pitch = data->attitude.pitch;
    zdrone_state.attitude.roll = data->attitude.roll;
    zdrone_state.attitude.yaw = data->attitude.yaw;
    zdrone_state.attitudeRate.roll = gx;
    zdrone_state.attitudeRate.pitch = gy;
    zdrone_state.attitudeRate.yaw = gz;

    // rt_printf_float("rMat[2][0]: ", imu->rMat[2][0], " \n");
    // acc_trace(zdrone_state.acc.x, zdrone_state.acc.y, zdrone_state.acc.z);
    // acc_trace(&fc_sensor_data, &zdrone_state, dt);
}

void Get_imu_data(fc_sensor_data_t *sensor, zdrone_state_t *state)
{
    state->attitude.pitch = sensor->attitude.pitch;
    state->attitude.roll = sensor->attitude.roll;
    state->attitude.yaw = sensor->attitude.yaw;

    state->attitudeRate.pitch = sensor->gyro.gyrolpf.y;
    state->attitudeRate.roll = sensor->gyro.gyrolpf.x;
    state->attitudeRate.yaw = sensor->gyro.gyrolpf.z;
}

float getCosTiltAngle(void)
{
    return fc_ctrl_imu.rMat[2][2];
}
